use std::any::Any;
use std::time::Duration;
use async_trait::async_trait;
use protoactor_rs::actor::{Actor, ActorContext, ActorError, DefaultActorContext, PID, Props, SupervisorDirective};
use protoactor_rs::supervision::{Strategies, RestartStatistics};
use protoactor_rs::supervision_event::SupervisionEventLogger;
use protoactor_rs::system::ActorSystem;
use tracing::{info, warn, debug, error, Level};

// Messages for our worker actor
enum WorkerMessage {
    DoWork,
    Fail,
    FailFatally,
}

// A worker actor that can fail
struct WorkerActor {
    name: String,
}

impl WorkerActor {
    fn new(name: String) -> Self {
        Self { name }
    }
}

#[async_trait]
impl Actor for WorkerActor {
    type Context = DefaultActorContext;
    
    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn Any + Send>) -> Result<(), ActorError> {
        if let Some(msg) = msg.downcast_ref::<WorkerMessage>() {
            match msg {
                WorkerMessage::DoWork => {
                    info!(actor = %self.name, "Worker doing work");
                    Ok(())
                }
                WorkerMessage::Fail => {
                    warn!(actor = %self.name, "Worker failing (recoverable)");
                    Err(ActorError::ActorFailure("Simulated failure".to_string()))
                }
                WorkerMessage::FailFatally => {
                    error!(actor = %self.name, "Worker failing fatally");
                    Err(ActorError::ActorPanic)
                }
            }
        } else {
            Ok(())
        }
    }
    
    async fn started(&self, _ctx: &mut Self::Context) -> Result<(), ActorError> {
        info!(actor = %self.name, "Worker started");
        Ok(())
    }
    
    async fn stopping(&self, _ctx: &mut Self::Context) -> Result<(), ActorError> {
        info!(actor = %self.name, "Worker stopping");
        Ok(())
    }
    
    async fn restarting(&self, _ctx: &mut Self::Context) -> Result<(), ActorError> {
        info!(actor = %self.name, "Worker restarting");
        Ok(())
    }
}

// Messages for our supervisor actor
enum SupervisorMessage {
    CreateWorker(String),
    TellWorkerToWork(String),
    TellWorkerToFail(String),
    TellWorkerToFailFatally(String),
}

// A supervisor actor that creates and manages worker actors
struct SupervisorActor {
    name: String,
    worker_refs: std::collections::HashMap<String, PID>,
}

impl SupervisorActor {
    fn new(name: String) -> Self {
        Self { 
            name,
            worker_refs: std::collections::HashMap::new(),
        }
    }
}

#[async_trait]
impl Actor for SupervisorActor {
    type Context = DefaultActorContext;
    
    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn Any + Send>) -> Result<(), ActorError> {
        let mut this = self.clone();
        
        if let Some(msg) = msg.downcast_ref::<SupervisorMessage>() {
            match msg {
                SupervisorMessage::CreateWorker(worker_name) => {
                    info!(supervisor = %this.name, worker = %worker_name, "Creating worker");
                    
                    // Create worker props with supervision strategy
                    let worker = WorkerActor::new(worker_name.clone());
                    let worker_props = Props::new()
                        .with_producer(Box::new(move || Box::new(worker.clone())))
                        .with_supervisor_strategy(
                            Strategies::one_for_one(
                                3, // Max retries
                                Duration::from_secs(10), // Within time window
                                |reason| match reason {
                                    ActorError::ActorPanic => SupervisorDirective::Stop,
                                    _ => SupervisorDirective::Restart,
                                }
                            )
                        );
                    
                    // Spawn the worker
                    let worker_pid = ctx.spawn(&worker_props).await?;
                    
                    // Store reference to the worker
                    this.worker_refs.insert(worker_name.clone(), worker_pid);
                }
                SupervisorMessage::TellWorkerToWork(worker_name) => {
                    if let Some(worker_pid) = this.worker_refs.get(worker_name) {
                        info!(supervisor = %this.name, worker = %worker_name, "Telling worker to work");
                        ctx.send(worker_pid, WorkerMessage::DoWork).await?;
                    }
                }
                SupervisorMessage::TellWorkerToFail(worker_name) => {
                    if let Some(worker_pid) = this.worker_refs.get(worker_name) {
                        info!(supervisor = %this.name, worker = %worker_name, "Telling worker to fail");
                        ctx.send(worker_pid, WorkerMessage::Fail).await?;
                    }
                }
                SupervisorMessage::TellWorkerToFailFatally(worker_name) => {
                    if let Some(worker_pid) = this.worker_refs.get(worker_name) {
                        info!(supervisor = %this.name, worker = %worker_name, "Telling worker to fail fatally");
                        ctx.send(worker_pid, WorkerMessage::FailFatally).await?;
                    }
                }
            }
        }
        
        Ok(())
    }
    
    async fn started(&self, _ctx: &mut Self::Context) -> Result<(), ActorError> {
        info!(supervisor = %self.name, "Supervisor started");
        Ok(())
    }
}

impl Clone for SupervisorActor {
    fn clone(&self) -> Self {
        Self {
            name: self.name.clone(),
            worker_refs: self.worker_refs.clone(),
        }
    }
}

// Custom decision function for our supervision strategy
fn custom_supervision_decider(reason: &ActorError) -> SupervisorDirective {
    match reason {
        ActorError::ActorPanic => {
            error!("Actor panicked, stopping");
            SupervisorDirective::Stop
        }
        ActorError::ActorFailure(msg) if msg.contains("fatal") => {
            error!("Fatal error, stopping: {}", msg);
            SupervisorDirective::Stop
        }
        ActorError::Timeout => {
            warn!("Timeout error, restarting");
            SupervisorDirective::Restart
        }
        _ => {
            info!("Other error, resuming: {:?}", reason);
            SupervisorDirective::Resume
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize the logger
    protoactor_rs::logging::init_global_logger(Some(Level::DEBUG), None)?;
    
    // Start the supervision event logger
    let _subscription = SupervisionEventLogger::start();
    
    // Create the actor system
    let system = ActorSystem::new()?;
    
    // Create a supervisor
    let supervisor = SupervisorActor::new("main-supervisor".to_string());
    let supervisor_props = Props::new()
        .with_producer(Box::new(move || Box::new(supervisor.clone())));
    
    let supervisor_pid = system.spawn(&supervisor_props).await?;
    
    // Create workers
    system.send(&supervisor_pid, SupervisorMessage::CreateWorker("worker1".to_string())).await?;
    system.send(&supervisor_pid, SupervisorMessage::CreateWorker("worker2".to_string())).await?;
    system.send(&supervisor_pid, SupervisorMessage::CreateWorker("worker3".to_string())).await?;
    
    // Give time for workers to be created
    tokio::time::sleep(Duration::from_millis(100)).await;
    
    // Send work messages
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToWork("worker1".to_string())).await?;
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToWork("worker2".to_string())).await?;
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToWork("worker3".to_string())).await?;
    
    // Wait a bit
    tokio::time::sleep(Duration::from_millis(100)).await;
    
    // Make workers fail (should be restarted)
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToFail("worker1".to_string())).await?;
    tokio::time::sleep(Duration::from_millis(100)).await;
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToFail("worker1".to_string())).await?;
    tokio::time::sleep(Duration::from_millis(100)).await;
    
    // Make sure worker1 still works after being restarted
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToWork("worker1".to_string())).await?;
    tokio::time::sleep(Duration::from_millis(100)).await;
    
    // Make worker2 fail fatally (should be stopped)
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToFailFatally("worker2".to_string())).await?;
    tokio::time::sleep(Duration::from_millis(100)).await;
    
    // Make worker3 fail multiple times beyond the retry limit
    for _ in 0..5 {
        system.send(&supervisor_pid, SupervisorMessage::TellWorkerToFail("worker3".to_string())).await?;
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
    
    // Wait to see all supervision actions take effect
    tokio::time::sleep(Duration::from_secs(1)).await;
    
    // Demonstrate that we can still use worker1
    system.send(&supervisor_pid, SupervisorMessage::TellWorkerToWork("worker1".to_string())).await?;
    
    // Clean up
    system.stop(&supervisor_pid).await?;
    
    Ok(())
} 